#include <iostream>
#include <string>
#include <iostream>
#include <memory>
#include <vector>
#include "compiler/lexer.h"
#include "compiler/parser.h"
#include "compiler/planner.h"
#include "execution/catalog.h"
#include "storage/disk_manager.h"
#include "storage/buffer_pool.h"

// Helper function to create and initialize necessary components
std::tuple<std::string, std::unique_ptr<Planner>, std::shared_ptr<Catalog>> setupTestEnvironment() {
    // Create necessary components
    std::string dbFile = "test_planner.db";
    auto diskManager = std::make_unique<DiskManager>(dbFile);
    auto bufferPool = std::make_unique<BufferPool>(10, std::move(diskManager));
    auto catalog = std::make_shared<Catalog>(std::move(bufferPool));
    
    // Initialize catalog
    catalog->loadFromStorage();
    
    // Create planner
    auto planner = std::make_unique<Planner>();
    
    return {dbFile, std::move(planner), catalog};
}

// Helper function to clean up test files
void cleanupTestEnvironment(const std::string& dbFile) {
    std::remove(dbFile.c_str());
    std::remove("catalog.dat");
}

// Test CREATE TABLE plan generation
void testCreateTablePlan() {
    std::cout << "Running CREATE TABLE plan generation test..." << std::endl;
    
    try {
        // Setup test environment
        auto [dbFile, planner, catalog] = setupTestEnvironment();
        
        // Test SQL statement
        std::string sql = "CREATE TABLE test_users (id INT, name VARCHAR, age INT)";
        
        // Tokenize and parse the SQL
        Lexer lexer(sql);
        std::vector<Token> tokens = lexer.tokenize();
        Parser parser(tokens);
        auto stmt = parser.parse();
        
        // Create execution plan
        std::unique_ptr<Execution::PlanNode> plan = planner->createPlan(stmt);
        
        // Verify the plan
        auto createTablePlan = dynamic_cast<Execution::CreateTablePlan*>(plan.get());
        if (!createTablePlan) {
            std::cerr << "Failed to create CreateTablePlan" << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        // Check plan details
        if (createTablePlan->tableName != "test_users") {
            std::cerr << "Incorrect table name: " << createTablePlan->tableName << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        if (createTablePlan->columns.size() != 3) {
            std::cerr << "Incorrect number of columns: " << createTablePlan->columns.size() << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        // Verify column details
        bool idColumnFound = false;
        bool nameColumnFound = false;
        bool ageColumnFound = false;
        
        for (const auto& [colName, colType] : createTablePlan->columns) {
            if (colName == "id" && colType == "INT") {
                idColumnFound = true;
            } else if (colName == "name" && colType == "VARCHAR") {
                nameColumnFound = true;
            } else if (colName == "age" && colType == "INT") {
                ageColumnFound = true;
            }
        }
        
        if (!idColumnFound || !nameColumnFound || !ageColumnFound) {
            std::cerr << "Missing or incorrect column definitions" << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        std::cout << "CREATE TABLE plan generation test passed!" << std::endl;
        std::cout << "Plan details: " << plan->toString() << std::endl;
        
        // Clean up
        cleanupTestEnvironment(dbFile);
        
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        exit(1);
    }
}

// Test INSERT plan generation
void testInsertPlan() {
    std::cout << "Running INSERT plan generation test..." << std::endl;
    
    try {
        // Setup test environment
        auto [dbFile, planner, catalog] = setupTestEnvironment();
        
        // Test SQL statement with explicit column list
        std::string sql = "INSERT INTO test_users (id, name, age) VALUES (1, 'Alice', 30)";
        
        // Tokenize and parse the SQL
        Lexer lexer(sql);
        std::vector<Token> tokens = lexer.tokenize();
        Parser parser(tokens);
        auto stmt = parser.parse();
        
        // Create execution plan
        std::unique_ptr<Execution::PlanNode> plan = planner->createPlan(stmt);
        
        // Verify the plan
        auto insertPlan = dynamic_cast<Execution::InsertPlan*>(plan.get());
        if (!insertPlan) {
            std::cerr << "Failed to create InsertPlan" << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        // Check plan details
        if (insertPlan->tableName != "test_users") {
            std::cerr << "Incorrect table name: " << insertPlan->tableName << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        if (insertPlan->columns.size() != 3) {
            std::cerr << "Incorrect number of columns: " << insertPlan->columns.size() << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        if (insertPlan->values.size() != 3) {
            std::cerr << "Incorrect number of values: " << insertPlan->values.size() << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        std::cout << "INSERT plan generation test passed!" << std::endl;
        std::cout << "Plan details: " << plan->toString() << std::endl;
        
        // Clean up
        cleanupTestEnvironment(dbFile);
        
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        exit(1);
    }
}

// Test INSERT plan generation without explicit column list
void testInsertPlanWithoutColumns() {
    std::cout << "Running INSERT plan generation without columns test..." << std::endl;
    
    try {
        // Setup test environment
        auto [dbFile, planner, catalog] = setupTestEnvironment();
        
        // Test SQL statement without explicit column list
        std::string sql = "INSERT INTO test_users VALUES (2, 'Bob', 25)";
        
        // Tokenize and parse the SQL
        Lexer lexer(sql);
        std::vector<Token> tokens = lexer.tokenize();
        Parser parser(tokens);
        auto stmt = parser.parse();
        
        // Create execution plan
        std::unique_ptr<Execution::PlanNode> plan = planner->createPlan(stmt);
        
        // Verify the plan
        auto insertPlan = dynamic_cast<Execution::InsertPlan*>(plan.get());
        if (!insertPlan) {
            std::cerr << "Failed to create InsertPlan" << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        // Check plan details
        if (insertPlan->tableName != "test_users") {
            std::cerr << "Incorrect table name: " << insertPlan->tableName << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        // Columns should be empty when not specified in SQL
        if (!insertPlan->columns.empty()) {
            std::cerr << "Columns should be empty when not specified: " << insertPlan->columns.size() << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        if (insertPlan->values.size() != 3) {
            std::cerr << "Incorrect number of values: " << insertPlan->values.size() << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        std::cout << "INSERT plan generation without columns test passed!" << std::endl;
        std::cout << "Plan details: " << plan->toString() << std::endl;
        
        // Clean up
        cleanupTestEnvironment(dbFile);
        
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        exit(1);
    }
}

// Test SELECT plan generation
void testSelectPlan() {
    std::cout << "Running SELECT plan generation test..." << std::endl;
    
    try {
        // Setup test environment
        auto [dbFile, planner, catalog] = setupTestEnvironment();
        
        // Test SQL statement with WHERE clause
        std::string sql = "SELECT id, name FROM test_users WHERE age > 25";
        
        // Tokenize and parse the SQL
        Lexer lexer(sql);
        std::vector<Token> tokens = lexer.tokenize();
        Parser parser(tokens);
        auto stmt = parser.parse();
        
        // Create execution plan
        std::unique_ptr<Execution::PlanNode> plan = planner->createPlan(stmt);
        
        // Verify the plan
        auto selectPlan = dynamic_cast<Execution::SelectPlan*>(plan.get());
        if (!selectPlan) {
            std::cerr << "Failed to create SelectPlan" << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        // Check plan details
        if (selectPlan->tableName != "test_users") {
            std::cerr << "Incorrect table name: " << selectPlan->tableName << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        if (selectPlan->columns.size() != 2) {
            std::cerr << "Incorrect number of columns: " << selectPlan->columns.size() << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        // Check specific columns
        bool idColumnFound = false;
        bool nameColumnFound = false;
        
        for (const auto& colName : selectPlan->columns) {
            if (colName == "id") {
                idColumnFound = true;
            } else if (colName == "name") {
                nameColumnFound = true;
            }
        }
        
        if (!idColumnFound || !nameColumnFound) {
            std::cerr << "Missing or incorrect selected columns" << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        // Check if filter is correctly generated
        if (selectPlan->filter.empty()) {
            std::cerr << "WHERE clause filter should not be empty" << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        std::cout << "SELECT plan generation test passed!" << std::endl;
        std::cout << "Plan details: " << plan->toString() << std::endl;
        
        // Clean up
        cleanupTestEnvironment(dbFile);
        
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        exit(1);
    }
}

// Test SELECT plan generation with * (all columns)
void testSelectPlanWithStar() {
    std::cout << "Running SELECT plan generation with * test..." << std::endl;
    
    try {
        // Setup test environment
        auto [dbFile, planner, catalog] = setupTestEnvironment();
        
        // Test SQL statement with * (select all columns)
        std::string sql = "SELECT * FROM test_users";
        
        // Tokenize and parse the SQL
        Lexer lexer(sql);
        std::vector<Token> tokens = lexer.tokenize();
        Parser parser(tokens);
        auto stmt = parser.parse();
        
        // Create execution plan
        std::unique_ptr<Execution::PlanNode> plan = planner->createPlan(stmt);
        
        // Verify the plan
        auto selectPlan = dynamic_cast<Execution::SelectPlan*>(plan.get());
        if (!selectPlan) {
            std::cerr << "Failed to create SelectPlan" << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        // Check plan details
        if (selectPlan->tableName != "test_users") {
            std::cerr << "Incorrect table name: " << selectPlan->tableName << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        // Check if * is correctly handled
        bool starFound = false;
        for (const auto& colName : selectPlan->columns) {
            if (colName == "*") {
                starFound = true;
                break;
            }
        }
        
        if (!starFound) {
            std::cerr << "* (select all) not correctly handled" << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        std::cout << "SELECT plan generation with * test passed!" << std::endl;
        std::cout << "Plan details: " << plan->toString() << std::endl;
        
        // Clean up
        cleanupTestEnvironment(dbFile);
        
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        exit(1);
    }
}

// Test DELETE plan generation
void testDeletePlan() {
    std::cout << "Running DELETE plan generation test..." << std::endl;
    
    try {
        // Setup test environment
        auto [dbFile, planner, catalog] = setupTestEnvironment();
        
        // Test SQL statement with WHERE clause
        std::string sql = "DELETE FROM test_users WHERE age < 18";
        
        // Tokenize and parse the SQL
        Lexer lexer(sql);
        std::vector<Token> tokens = lexer.tokenize();
        Parser parser(tokens);
        auto stmt = parser.parse();
        
        // Create execution plan
        std::unique_ptr<Execution::PlanNode> plan = planner->createPlan(stmt);
        
        // Verify the plan
        auto deletePlan = dynamic_cast<Execution::DeletePlan*>(plan.get());
        if (!deletePlan) {
            std::cerr << "Failed to create DeletePlan" << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        // Check plan details
        if (deletePlan->tableName != "test_users") {
            std::cerr << "Incorrect table name: " << deletePlan->tableName << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        // Check if filter is correctly generated
        if (deletePlan->filter.empty()) {
            std::cerr << "WHERE clause filter should not be empty" << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        std::cout << "DELETE plan generation test passed!" << std::endl;
        std::cout << "Plan details: " << plan->toString() << std::endl;
        
        // Clean up
        cleanupTestEnvironment(dbFile);
        
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        exit(1);
    }
}

// Test DELETE plan generation without WHERE clause
void testDeletePlanWithoutWhere() {
    std::cout << "Running DELETE plan generation without WHERE clause test..." << std::endl;
    
    try {
        // Setup test environment
        auto [dbFile, planner, catalog] = setupTestEnvironment();
        
        // Test SQL statement without WHERE clause (delete all rows)
        std::string sql = "DELETE FROM test_users";
        
        // Tokenize and parse the SQL
        Lexer lexer(sql);
        std::vector<Token> tokens = lexer.tokenize();
        Parser parser(tokens);
        auto stmt = parser.parse();
        
        // Create execution plan
        std::unique_ptr<Execution::PlanNode> plan = planner->createPlan(stmt);
        
        // Verify the plan
        auto deletePlan = dynamic_cast<Execution::DeletePlan*>(plan.get());
        if (!deletePlan) {
            std::cerr << "Failed to create DeletePlan" << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        // Check plan details
        if (deletePlan->tableName != "test_users") {
            std::cerr << "Incorrect table name: " << deletePlan->tableName << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        // Filter should be empty when no WHERE clause is specified
        if (!deletePlan->filter.empty()) {
            std::cerr << "Filter should be empty when no WHERE clause is specified" << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
        }
        
        std::cout << "DELETE plan generation without WHERE clause test passed!" << std::endl;
        std::cout << "Plan details: " << plan->toString() << std::endl;
        
        // Clean up
        cleanupTestEnvironment(dbFile);
        
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        exit(1);
    }
}

// Test planner error handling
void testPlannerErrorHandling() {
    std::cout << "Running planner error handling test..." << std::endl;
    
    try {
        // Setup test environment
        auto [dbFile, planner, catalog] = setupTestEnvironment();
        
        // Test with invalid SQL that can be parsed but not planned
        // This is a hypothetical example, actual error behavior may vary
        std::string sql = "CREATE TABLE (id INT)";
        
        try {
            // Tokenize and parse the SQL
            Lexer lexer(sql);
            std::vector<Token> tokens = lexer.tokenize();
            Parser parser(tokens);
            auto stmt = parser.parse();
            
            // This should throw an exception
            planner->createPlan(stmt);
            
            // If we get here, the error wasn't caught
            std::cerr << "Planner did not catch invalid statement" << std::endl;
            cleanupTestEnvironment(dbFile);
            exit(1);
            
        } catch (const std::exception& e) {
            // This is expected
            std::cout << "Expected planner error caught: " << e.what() << std::endl;
        }
        
        std::cout << "Planner error handling test passed!" << std::endl;
        
        // Clean up
        cleanupTestEnvironment(dbFile);
        
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        exit(1);
    }
}

int main() {
    std::cout << "Running Planner unit tests..." << std::endl;
    std::cout << "==================================" << std::endl;
    
    // Run individual tests
    testCreateTablePlan();
    testInsertPlan();
    testInsertPlanWithoutColumns();
    testSelectPlan();
    testSelectPlanWithStar();
    testDeletePlan();
    testDeletePlanWithoutWhere();
    testPlannerErrorHandling();
    
    std::cout << "==================================" << std::endl;
    std::cout << "All Planner tests passed successfully!" << std::endl;
    return 0;
}